package com.tianji.aigc.memory;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.tianji.aigc.entity.DBChatMemory;
import com.tianji.aigc.service.ChatMemoryService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.model.Media;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @author wys
 * @version V1.0
 * @date 2025-04-23 0:09
 */
public class SQLChatMemory implements ChatMemory {
    @Resource
    private ChatMemoryService chatMemoryService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        if (CollUtil.isEmpty(messages)) {
            // 如果消息列表为空则直接返回
            return;
        }
        messages.forEach(message -> {
            DBChatMemory dbChatMemory = new DBChatMemory();
            dbChatMemory.setSessionId(conversationId);
            dbChatMemory.setMessageType(message.getMessageType().name());
            dbChatMemory.setMetadata(JSONUtil.toJsonStr(message.getMetadata()));
            dbChatMemory.setTextContent(message.getText());
            if (message instanceof UserMessage userMessage) {
                dbChatMemory.setMedia(JSONUtil.toJsonStr(userMessage.getMedia()));
            }
            if (message instanceof AssistantMessage assistantMessage) {
                dbChatMemory.setToolCalls(JSONUtil.toJsonStr(assistantMessage.getToolCalls()));
                dbChatMemory.setMedia(JSONUtil.toJsonStr(assistantMessage.getMedia()));
            }
            if (message instanceof ToolResponseMessage toolResponseMessage) {
                dbChatMemory.setToolResponses(JSONUtil.toJsonStr(toolResponseMessage.getResponses()));
            }
            dbChatMemory.setCreateTime(LocalDateTime.now());
            dbChatMemory.setUpdateTime(LocalDateTime.now());
            chatMemoryService.save(dbChatMemory);
        });
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<DBChatMemory> dbChatMemories = chatMemoryService.list(new LambdaQueryWrapper<DBChatMemory>()
                .eq(DBChatMemory::getSessionId, conversationId)
                .last("limit " + lastN)
                .orderByAsc(DBChatMemory::getUpdateTime));
        List<Message> messages = new ArrayList<>();
        for (DBChatMemory dbChatMemory : dbChatMemories) {
            MessageType messageType = MessageType.valueOf(dbChatMemory.getMessageType());
            switch (messageType) {
                case SYSTEM -> {
                    messages.add(new SystemMessage(dbChatMemory.getTextContent()));
                }
                case USER -> {
                    messages.add(new UserMessage(dbChatMemory.getTextContent(),
                            JSONUtil.toList(dbChatMemory.getMedia(), Media.class),
                            JSONUtil.toBean(dbChatMemory.getMetadata(), Map.class)));
                }
                case ASSISTANT -> {
                    messages.add(new AssistantMessage(dbChatMemory.getTextContent()
                            , JSONUtil.toBean(dbChatMemory.getProperties(), Map.class)
                            , JSONUtil.toList(dbChatMemory.getToolCalls(), AssistantMessage.ToolCall.class)
                            , JSONUtil.toList(dbChatMemory.getMedia(), Media.class)));
                }
                case TOOL -> {
                    messages.add(new ToolResponseMessage(
                            JSONUtil.toList(dbChatMemory.getToolResponses(), ToolResponseMessage.ToolResponse.class),
                            JSONUtil.toBean(dbChatMemory.getMetadata(), Map.class)));
                }
            }
        }
        return messages;
    }

    @Override
    public void clear(String conversationId) {
        chatMemoryService.remove(new LambdaQueryWrapper<DBChatMemory>().eq(DBChatMemory::getSessionId, conversationId));
    }
}
